from llvmlite import binding as llvm
from ctypes import CFUNCTYPE, c_int32, c_double, POINTER, c_void_p, cast

# 初始化 LLVM
llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter()

# 定义 LLVM 汇编代码
ir_text = """
; 定义函数指针类型
%func_ptr = type i32 (i32, i32)*
%math_func = type double (double, double)*

; 基本数学运算函数
define i32 @add(i32 %a, i32 %b) {
  %sum = add i32 %a, %b
  ret i32 %sum
}

define i32 @sub(i32 %a, i32 %b) {
  %diff = sub i32 %a, %b
  ret i32 %diff
}

define double @power(double %a, double %b) {
  %result = call double @llvm.pow.f64(double %a, double %b)
  ret double %result
}

; 使用函数指针作为参数的函数
define i32 @apply_op(%func_ptr %op, i32 %a, i32 %b) {
  %result = call i32 %op(i32 %a, i32 %b)
  ret i32 %result
}

; 返回函数指针的函数
define %func_ptr @get_operation(i1 %is_add) {
  %op = select i1 %is_add, %func_ptr @add, %func_ptr @sub
  ret %func_ptr %op
}

; 高阶函数：接受函数指针并返回新函数指针
define %math_func @create_math_op(%math_func %base_op, double %factor) {
  ; 分配内存存储因子
  %factor_ptr = alloca double
  store double %factor, double* %factor_ptr
  
  ; 返回一个新函数
  ret %math_func @apply_factor
}

; 应用因子的函数
define double @apply_factor(double %a, double %b) {
  ; 从闭包中获取因子
  %factor_ptr = bitcast i8* null to double*
  %factor = load double, double* %factor_ptr
  
  ; 调用基础操作
  %base_result = call double @power(double %a, double %b)
  
  ; 应用因子
  %result = fmul double %base_result, %factor
  ret double %result
}

; 函数指针回调示例
define double @callback_example(%math_func %callback, double %a, double %b) {
  %result = call double %callback(double %a, double %b)
  ret double %result
}

; 声明外部函数
declare double @llvm.pow.f64(double, double)
"""

# 解析汇编代码
module = llvm.parse_assembly(ir_text)
module.verify()

# 创建执行引擎
target = llvm.Target.from_default_triple()
target_machine = target.create_target_machine()
engine = llvm.create_mcjit_compiler(module, target_machine)

# 获取LLVM函数地址
add_addr = engine.get_function_address("add")
sub_addr = engine.get_function_address("sub")
power_addr = engine.get_function_address("power")
apply_op_addr = engine.get_function_address("apply_op")
get_operation_addr = engine.get_function_address("get_operation")
callback_example_addr = engine.get_function_address("callback_example")

# 定义函数类型
add_func_type = CFUNCTYPE(c_int32, c_int32, c_int32)
sub_func_type = CFUNCTYPE(c_int32, c_int32, c_int32)
power_func_type = CFUNCTYPE(c_double, c_double, c_double)
apply_op_type = CFUNCTYPE(c_int32, c_void_p, c_int32, c_int32)
get_operation_type = CFUNCTYPE(c_void_p, c_int32)
callback_example_type = CFUNCTYPE(c_double, c_void_p, c_double, c_double)

# 创建函数对象
add_func = add_func_type(add_addr)
sub_func = sub_func_type(sub_addr)
power_func = power_func_type(power_addr)
apply_op_func = apply_op_type(apply_op_addr)
get_operation_func = get_operation_type(get_operation_addr)
callback_example_func = callback_example_type(callback_example_addr)

# 定义Python端的数学函数
def py_multiply(a: float, b: float) -> float:
    return a * b

def py_divide(a: float, b: float) -> float:
    return a / b if b != 0 else float('nan')

# 将Python函数转换为函数指针
py_multiply_type = CFUNCTYPE(c_double, c_double, c_double)
py_multiply_func = py_multiply_type(py_multiply)
py_multiply_ptr = cast(py_multiply_func, c_void_p).value

py_divide_type = CFUNCTYPE(c_double, c_double, c_double)
py_divide_func = py_divide_type(py_divide)
py_divide_ptr = cast(py_divide_func, c_void_p).value

# 测试函数调用
if __name__ == "__main__":
    print("=== 基本函数调用 ===")
    print(f"add(5, 3) = {add_func(5, 3)}")
    print(f"sub(5, 3) = {sub_func(5, 3)}")
    print(f"power(2.0, 3.0) = {power_func(2.0, 3.0):.1f}")
    
    print("\n=== 函数指针作为参数 ===")
    # 将add函数转换为void指针
    add_ptr = cast(add_func, c_void_p).value
    sub_ptr = cast(sub_func, c_void_p).value
    
    print(f"apply_op(add, 5, 3) = {apply_op_func(add_ptr, 5, 3)}")
    print(f"apply_op(sub, 5, 3) = {apply_op_func(sub_ptr, 5, 3)}")
    
    print("\n=== 返回函数指针 ===")
    # 获取函数指针 (1 表示 true，获取 add 函数)
    add_ptr_ret = get_operation_func(1)
    # 使用返回的函数指针
    add_func_ret = cast(add_ptr_ret, add_func_type)
    print(f"get_operation(true) -> add(7, 2) = {add_func_ret(7, 2)}")
    
    # 获取 sub 函数指针 (0 表示 false)
    sub_ptr_ret = get_operation_func(0)
    sub_func_ret = cast(sub_ptr_ret, sub_func_type)
    print(f"get_operation(false) -> sub(7, 2) = {sub_func_ret(7, 2)}")
    
    print("\n=== 回调Python函数 ===")
    # 使用Python的乘法函数
    print(f"callback_example(py_multiply, 4.0, 5.0) = {callback_example_func(py_multiply_ptr, 4.0, 5.0):.1f}")
    
    # 使用Python的除法函数
    print(f"callback_example(py_divide, 10.0, 2.0) = {callback_example_func(py_divide_ptr, 10.0, 2.0):.1f}")
    
    # 使用LLVM的power函数
    print(f"callback_example(power, 3.0, 2.0) = {callback_example_func(power_addr, 3.0, 2.0):.1f}")
    
    print("\n=== 闭包函数指针 (模拟) ===")
    # 注意：我们的IR中闭包实现是简化的，实际需要更复杂的内存管理
    # 这里我们只演示概念
    print("(示例中闭包实现需要更完整的内存管理)")
